#pragma once
#include <map>
#include <Eigen/Eigen>
#include <utility>
#include "HairSimulator.h"
#include "HairReader.h"
#include "SDF.h"

namespace XRwy {

    struct MassSpring {
        //local indicates whether we should substract the strand index when we apply to vector or matrixs
        virtual void apply_vector(Eigen::VectorXf &vec, bool local) const = 0;
        virtual void apply_triplet(std::vector<Eigen::Triplet<float>> &triplet, float t, bool local) const = 0;
        virtual ~MassSpring() = default;
    };

    class HairMassSpringSimulator;

    struct NormalMassSpring: public MassSpring {
        int idx1, idx2, strand_idx;
        float k;
        float l0;
        float b;
        float k_div_l0;
        HairMassSpringSimulator *simulator;

        using SegmentType = Eigen::VectorXf::SegmentReturnType;

        NormalMassSpring(HairMassSpringSimulator *simulator_, int idx1_, int idx2_, int strand_idx_, float k_, float l0_);
        NormalMassSpring(HairMassSpringSimulator *simulator_, int idx1_, int idx2_, int strand_idx_, float k_);

        virtual void apply_vector(Eigen::VectorXf &vec, bool local) const;

        virtual void apply_triplet(std::vector<Eigen::Triplet<float>> &triplet, float t, bool local) const;

        virtual ~NormalMassSpring() = default;

        SegmentType get_pos(int idx) const;
        SegmentType get_vel(int idx) const;
        SegmentType get_rest_pos(int idx) const;

        inline SegmentType get_pos1() const { return get_pos(idx1); }
        inline SegmentType get_pos2() const { return get_pos(idx2); }
        inline SegmentType get_vel1() const { return get_vel(idx1); }
        inline SegmentType get_vel2() const { return get_vel(idx2); }
        inline SegmentType get_rest_pos1() const { return get_rest_pos(idx1); }
        inline SegmentType get_rest_pos2() const { return get_rest_pos(idx2); }
    };

    enum HairParticleType {
		NORMAL,
		VIRTUAL
	};

	struct StrandInfo {
		int index;
		int visible_particle_count;
		std::vector<HairParticleType> particle_types;
        std::vector<MassSpring *> springs;

		StrandInfo(): index(0), visible_particle_count(0), particle_types(std::vector<HairParticleType>()) {}

		inline int get_particle_count() const { return particle_types.size(); }
		inline int get_visible_particle_count() const { return visible_particle_count; }
	};

	struct ParticleInfo {
		int strand_index;
		Eigen::Vector3f rest_pos;

		ParticleInfo(int strand_index_, Eigen::Vector3f rest_pos_) :  strand_index(strand_index_), rest_pos(rest_pos_) {}
	};

	struct HairMassSpringConfiguration {
		bool enable_body_collision;
        bool enable_hair_stiction;
		float ke; //k for "edge spring" and "extra edge spring"
		float kb; //k for "bending spring"
		float kt; //k for "torsion spring"
		float kr; //k for "self-replusion"
		float b; //b for "damping"
		float npass; //number of pass for iteration
		float mass; //particle mass
		float dr; //distance for "stiction" to occur
        float invalid_dr; //the max distance for the "stiction" spring
		float max_replusion_link; //max spring count for every particle
		Eigen::Vector3f extra_force; //the force like gravity
        Eigen::Affine3f init_transformation; //the initial transformation
	};

	class HairMassSpringSimulator : public HairSimulator {
		friend class NormalMassSpring;
	public:
		HairMassSpringSimulator(const HairReader &hair_reader, const HairMassSpringConfiguration &conf);

		virtual int get_particle_count() const;
		virtual int get_strand_count() const;
		virtual int get_particle_count_on_strand(int idx) const;
		virtual void copy_particles(float *positions) const;

		std::vector<std::pair<HairParticleType, Eigen::Vector3f>> dump_all_particles_info() const;

		virtual void on_frame(const HairSimulatorData &data);

		virtual ~HairMassSpringSimulator();
	private:
		Eigen::VectorXf rest_pos, pos, vel, filter;

		int particle_count;
		int vector_dim;
		int visible_particle_count;
		std::vector<StrandInfo> strand_infos;
		std::vector<NormalMassSpring> stiction_springs;
		std::vector<ParticleInfo> particle_infos;

		std::vector<int> aligned_x_particle_indexes; //sorted by x-axis

		HairMassSpringConfiguration conf;

		void update_aligned_hierarchy();

        void resolve_collision(const HairSimulatorData &data, float t);
		void resolve_time_integration(const HairSimulatorData &data, const std::vector<Eigen::Vector3f> strand_root_vels, float t);
        void resolve_stiction(const HairSimulatorData &data, float t);
        void update_stiction_spring();
	};
}